import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
from torchvision import transforms
import matplotlib.pyplot as plt
import numpy as np
import os
import time
from sklearn.metrics import classification_report

# 超参数
BATCH_SIZE = 64
EPOCHS = 10
LR = 0.001
MODEL_NAME = "gru_mlp"

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建保存目录
os.makedirs('models', exist_ok=True)
os.makedirs('visualization', exist_ok=True)
os.makedirs('results', exist_ok=True)

# 数据预处理
transform = transforms.Compose([
    transforms.Resize((32, 32)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

# 加载数据集
train_dataset = torchvision.datasets.MNIST(root='./data',
                                           train=True,
                                           download=True,
                                           transform=transform)
test_dataset = torchvision.datasets.MNIST(root='./data',
                                          train=False,
                                          transform=transform)

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                           batch_size=BATCH_SIZE,
                                           shuffle=True)
test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                          batch_size=BATCH_SIZE,
                                          shuffle=False)


# GRU - MLP 模型定义
class GRU_MLP(nn.Module):
    def __init__(self):
        super(GRU_MLP, self).__init__()
        self.gru = nn.GRU(input_size=32, hidden_size=128, num_layers=1, batch_first=True)
        self.fc1 = nn.Linear(128, 64)
        self.fc2 = nn.Linear(64, 10)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = x.squeeze(1)  # 去除通道维度
        out, _ = self.gru(x)
        out = out[:, -1, :]  # 取最后一个时间步的输出
        out = self.relu(self.fc1(out))
        out = self.fc2(out)
        return out


model = GRU_MLP().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=LR)

# 训练记录
train_losses = []
train_accs = []
test_losses = []
test_accs = []
start_time = time.time()
best_acc = 0.0

# 训练循环
for epoch in range(EPOCHS):
    epoch_start = time.time()
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0

    for images, labels in train_loader:
        images = images.to(device)
        labels = labels.to(device)

        outputs = model(images)
        loss = criterion(outputs, labels)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()

    train_loss = running_loss / len(train_loader)
    train_acc = correct / total
    train_losses.append(train_loss)
    train_accs.append(train_acc)

    # 测试阶段
    model.eval()
    test_loss = 0.0
    correct = 0
    total = 0
    all_preds = []
    all_labels = []

    with torch.no_grad():
        for images, labels in test_loader:
            images = images.to(device)
            labels = labels.to(device)

            outputs = model(images)
            loss = criterion(outputs, labels)

            test_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()

            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())

    test_loss = test_loss / len(test_loader)
    test_acc = correct / total
    test_losses.append(test_loss)
    test_accs.append(test_acc)

    # 保存最佳模型
    if test_acc > best_acc:
        best_acc = test_acc
        torch.save(model.state_dict(), f'models/{MODEL_NAME}_best.pth')

    epoch_time = time.time() - epoch_start
    print(f'Epoch [{epoch + 1}/{EPOCHS}], '
          f'Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, '
          f'Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}, '
          f'Time: {epoch_time:.2f}s')

total_time = time.time() - start_time

# 保存最终模型
torch.save(model.state_dict(), f'models/{MODEL_NAME}_final.pth')

# 绘制曲线
plt.figure(figsize=(12, 4))
plt.subplot(1, 3, 1)
plt.plot(train_losses, label='Train')
plt.plot(test_losses, label='Test')
plt.title('Loss Curve')
plt.legend()

plt.subplot(1, 3, 2)
plt.plot(train_accs, label='Train')
plt.plot(test_accs, label='Test')
plt.title('Accuracy Curve')
plt.legend()

plt.subplot(1, 3, 3)
plt.plot(1 - np.array(test_accs), label='Test')
plt.title('Misclassification Rate Curve')
plt.legend()

plt.tight_layout()
plt.savefig(f'visualization/{MODEL_NAME}_curves.png')
plt.close()

# 生成分类报告
report = classification_report(all_labels, all_preds, digits=4)
with open(f'results/{MODEL_NAME}_report.txt', 'w') as f:
    f.write(report)

# 保存训练信息
with open(f'results/{MODEL_NAME}_metrics.txt', 'w') as f:
    f.write(f'Training Time: {total_time:.2f} seconds\n')
    if torch.cuda.is_available():
        f.write(f'Max GPU Memory Used: {torch.cuda.max_memory_allocated() / 1024 ** 2:.2f} MB\n')
    f.write(f'Final Test Accuracy: {test_accs[-1]:.4f}\n')
    f.write(f'Best Test Accuracy: {best_acc:.4f}\n')

# 可视化预测示例
model.eval()
with torch.no_grad():
    dataiter = iter(test_loader)
    images, labels = next(dataiter)
    images = images.to(device)
    outputs = model(images)
    _, preds = torch.max(outputs, 1)

fig = plt.figure(figsize=(8, 8))
for idx in range(9):
    ax = fig.add_subplot(3, 3, idx + 1)
    img = images[idx].cpu().numpy().squeeze()
    ax.imshow(img, cmap='gray')
    ax.set_title(f'Pred: {preds[idx]} | True: {labels[idx]}')
    ax.axis('off')
plt.tight_layout()
plt.savefig(f'visualization/{MODEL_NAME}_predictions.png')
plt.close()

print("Training completed and all results saved!")